/**
 CollinsKt.cxx
 
 \file
 Implementation of class CollinsKt.
 
 \author Thomas Burton 
 \date 2/28/12
 \copyright 2012 BNL. All rights reserved.
 */

#include <iostream>
#include <map>
#include <sstream>
#include <utility>

#include <boost/foreach.hpp>

#include <Math/ParamFunctor.h>
#include <TCanvas.h>
#include <TH1D.h>
#include <TMath.h>
#include <TRandom.h>

#include "tppmc/Bins.h"
#include "tppmc/CollinsKt.h"
#include "tppmc/DssFragmentation.h"
#include "tppmc/KeyedLists.h"
#include "tppmc/KkpParton.h"
#include "tppmc/parameters.h"

namespace tppmc {
   
   namespace collins {

      /**
       Function object evaluating Collins function.
       */
      struct FunctorNew {
         
         // <pT^2> for the Gaussian pT distribution
         double mMeanPt2;
         // The unpolarised fragmentation function against which the positivity
         // bound is computed
         boost::shared_ptr<tppmc::FragmentationFunction> mUnpol;

         FunctorNew(boost::shared_ptr<tppmc::FragmentationFunction> unpol)
         : mMeanPt2(0.2)
         , mUnpol(unpol) {
         }

         enum EFunctorParameters {
            SCALE = 0,
            FAVOURED = 1,
            POLARISATION = 2
         };

         double Favoured(const Parton& parton, const Hadron& hadron,
                         double Q2) {
            return mUnpol->Evaluate(parton, hadron, Q2);
         }
         
         double Unfavoured(const Parton& parton, const Hadron& hadron,
                           double Q2) {
            return -Favoured(parton, hadron, Q2);
         }

         double operator()(double* point, double* parameters) {
            double z = point[0];
            double phi = point[1];
            double qs2 = pow(parameters[SCALE], 2.);
            int isFavoured = static_cast<int>(parameters[FAVOURED]);
            double depol_factor = parameters[POLARISATION];
            // The fragmentation functions are parametrized in terms of the
            // fragmentation function for a positive pion, and we choose the
            // favored or unfavored fragmentation function from this by
            // setting the quark type, up quark for favoured,
            // or down quark for unfavoured.
            double D1(0.);
            double collins(0.);
            // Favoured collins = D1(u-->pi+)
            // Unfavoured collins = -D1(d-->pi+) (note the minus sign)
            // For pi0, use average of favoured and unfavoured for that quark
            Hadron piplus;
            piplus.type = 211;
            piplus.z = z;
            if(0 == isFavoured) { // unfavoured
               D1 = mUnpol->Evaluate(tppmc::Parton(0., 0., 0., 1), piplus, qs2);
               collins = -D1;
            } // if
            else if(1 == isFavoured) {// favoured
               D1 = mUnpol->Evaluate(tppmc::Parton(0., 0., 0., 2), piplus, qs2);
               collins = D1;
            } // else if
            else if(2 == isFavoured) { // Mean of favoured and unfavoured
               // Based this on the DSS definition of pi0 fragmentation.
               // They use (q->pi0) = average of (q->pi+) and (q->pi-)
               // which is the same as (favoured->pi+) and (unfavoured->pi-)
               // with the charge conjugation assumptions.
               // e.g. their u->pi- is the same as d->pi+
               double favoured = mUnpol->Evaluate(tppmc::Parton(0., 0., 0., 2), piplus, qs2);
               double unfavoured = mUnpol->Evaluate(tppmc::Parton(0., 0., 0., 1), piplus, qs2);
               // Unpolarised is the average
               D1 = 0.5 * (favoured + unfavoured);
               // Unfavoured collins is negative of D1
               collins = 0.5 * (favoured - unfavoured);
            } // else if
            // Why is there no "spin" factor (sign) in this equation? Because what we want is the angular 
            // distribution referenced to the spin direction. This is accounted for properly 
            // when we calculate the angular distrubution of the hadrons in the fragmentation
            // because it is done with respect to the spin direction.
            // NOTE ALSO: The angular distribution that is given is with respect to the vector
            // p^ x jT^, so the dot product gives a cos(phi) behavior. Because this is perp. to
            // jt^ the hadron distribution will have a sin(phi) modulation.
            double full_FF = D1 + cos(phi) * depol_factor * collins;
            return std::max(full_FF, 0.); // Don't go negative!
         }
      }; // FunctorNew


      //
      // class PositivityKt
      //


      void PositivityKt::Qa(std::string filename) {
         TCanvas canvas;
         canvas.Print((filename + "[").c_str());
         double depols[] = { 0.1 };
         UniformBins depolbins(1, 0.1, 0.9);
         mFunctions.Initialise(&depolbins, mQBins);
         std::ostringstream oss;
         std::cout << Partons().size() << " " << Hadrons().size() << std::endl;
         BOOST_FOREACH(int parton, Partons()) {
            std::cout << parton << std::endl;
            BOOST_FOREACH(int hadron, Hadrons()) {
               BOOST_FOREACH(double depol, depols) {
                  Hadron thehadron;
                  thehadron.type = hadron;
                  thehadron.mParton.type = parton;
                  thehadron.mParton.cos_spin_depol = depol;
                  TF2* func = mFunctions.Find(thehadron, 2.);
                  if(not func) {
                     std::cout << "Buh?" << std::endl;
                     continue;
                  } // if
                  oss << "Parton " << parton << " hadron " << hadron
                  << " depol " << depol;
                  std::cout << oss.str() << std::endl;
                  TH1D hist("", oss.str().c_str(), 9, 0.05, 0.95);
                  oss.str("");
                  for(int i(1); i <= 9; ++i) {
                     double integral = func->Integral(0.05, 0.05 + i * 0.1,
                        -TMath::PiOver2(), TMath::PiOver2());
                     hist.SetBinContent(i, integral);
                  } // for
                  hist.Scale(1. / hist.GetBinContent(hist.GetNbinsX()));
                  hist.SetMinimum(0.);
                  hist.Draw();
                  canvas.Print(filename.c_str());
               } // BOOST_FOREACH
            } // BOOST_FOREACH
         } // BOOST_FOREACH
         canvas.Print((filename + "]").c_str());
      }

      PositivityKt::PositivityKt()
      : mMeanPt2(0.2) {
      }

      PositivityKt::PositivityKt(FragFcnPtr unpol, double meanPt2)
      : mMeanPt2(meanPt2)
      , mUnpol(unpol) {
      }

      bool PositivityKt::Initialise() {
         if(not mUnpol.get()) {
            mUnpol = FragFcnPtr(new DssFragmentation(LO));
         } // if
         mQBins = new UniformBins(230, sqrt(MinimumScale()), sqrt(MaximumScale()));
         UniformBins depol(20, 0., 1.);
         mFunctions.Initialise(&depol, mQBins);
         int hadrons[3] = { -211, 111, 211 };
         mHadrons.assign(hadrons, hadrons + 3);
         int partons[6] = { 2, -2, 1, -1, 3, -3 };
         mPartons.assign(partons, partons + 6);
         return mQBins not_eq NULL;
      }

      const std::vector<int>& PositivityKt::Hadrons() const {
         return mHadrons;
      }

      const std::vector<int>& PositivityKt::Partons() const {
         return mPartons;
      }

      PositivityKt::~PositivityKt() { // Nada - only smart pointers are used
      }

      void PositivityKt::ComputeFragmentationProbabilities() {
         std::vector<double> QBins = mQBins->Edges();
         std::vector<int> hadronlist = Hadrons();
         std::vector<int> partonlist = Partons();
         BOOST_FOREACH(int parton, partonlist) {
            FragTablePtr table(new FragmentationTable(mQBins->Edges(),
                                                      hadronlist));
            mProbabilities.insert(std::make_pair(parton, table));
            // Loop over hadrons and Q bins and integrate the function
            BOOST_FOREACH(int hadron, hadronlist) {
               Hadron part;
               part.type = hadron;
               part.mParton.type = parton;
               // Overall probabilities of hadron production don't depend on
               // polarisation, so use zero
               part.mParton.cos_spin_depol = 0.;
               for(int i(0); i < mQBins->Number(); ++i) {
                  double Q = QBins.at(i); // Lower edge of bin
                  TF2* function = mFunctions.Find(part, Q);
                  assert(function);
                  double integral = function->Integral(ZMin(), ZMax(),
                     -TMath::Pi(), TMath::Pi());
                  mProbabilities[parton]->Set(hadron, Q, integral);
               } // for
            } // BOOST_FOREACH
            mProbabilities[parton]->Normalise(1.);
         } // BOOST_FOREACH
         
      }

      int PositivityKt::ThrowHadronType(const Parton& parton, double Q2) {
         if(INVALID_HADRON not_eq mHadron) { // Force to pick certain type
            return mHadron;
         } // if
         if(mProbabilities.empty()) { // Lazy initialisation
            ComputeFragmentationProbabilities();
         } // if
         FragTablePtr table = mProbabilities[parton.type];
         /**\todo Need protection against invalid result */
         // Passed Q2 but table stores probabilities in bins in Q --> sqrt
         // Don't allow it to go below the minimum scale
         double Q = std::max(sqrt(Q2), sqrt(mUnpol->MinimumScale()));
         std::pair<int, bool> hadron = table->Random(Q);
         if(not hadron.second) {
            std::cerr << "Threw invalid hadron type... using pi0" << std::endl;
            hadron.first = 111; // PIZERO;
         } // if
         return hadron.first;
      }

      std::vector<Hadron> PositivityKt::Fragment(const Parton& parton,
                                                 double scale) {
         if(not Supports(parton)) {
            return mUnpol->Fragment(parton, scale);
         } // if
         Hadron hadron;
         hadron.type = ThrowHadronType(parton, scale);
         switch(hadron.type) {
            case  211:
            case -211: // Deliberate fall-through
            case  111: // Deliberate fall-through
               break;
            default:
               std::cerr << "Threw hadron type " << hadron.type << " for parton " << 
               parton << " at scale " << scale << "? Returning DSS instead" << 
               std::endl;
               return mUnpol->Fragment(parton, scale);
         } // switch(h)
         hadron.mParton = parton;
         // Find the function with the parameters for this hadron/parton/scale
         // Input scale is Q2 but functions are stored in bins in Q
         TF2* function = mFunctions.Find(hadron, sqrt(scale));
         hadron.favoured = function->GetParameter(FunctorNew::FAVOURED);
         if(not function) {
            std::cerr << "No function for hadron " << hadron.type <<
            " parton " << hadron.mParton.type << " depol " <<
            hadron.mParton.cos_spin_depol << " at scale " << sqrt(scale)
            <<std::endl;
            assert(function);
         } // if
         // Randomly throw z and phi according to the function
         // Don't use TF2::SetRange() to restrict the z range to throw, as this
         // clears the integral cache. Instead keep throwing until until we
         // get a value in the required range.
         hadron.z = -1.;
         while(hadron.z > mZUpper or hadron.z < mZLower) {
            function->GetRandom2(hadron.z, hadron.phigen);
         } // while
         // Randomly generate kt.
         double jt_sigma = sqrt(mMeanPt2/2.0);  
         double jt; 
         // Throw jt according to an exponential distribution such
         // that mean jt^2 is 0.2
         do {
            jt = jt_sigma*sqrt( -2.0*log(gRandom->Uniform(1.0)) ); 	  
         } while(jt>(3.0*jt_sigma));
         TVector3 mom(jt, 0., 0.);
         // TVector3::SetPhi() rotates phi,
         // leaving magnitude and theta unchanged.
         mom.SetPhi(hadron.phigen);
         hadron.px = mom.Px();
         hadron.py = mom.Py();
         hadron.pz = parton.PxPyPz().Mag() * hadron.z;
         hadron.mFrag = Evaluate(parton, hadron, scale);
         hadron.mCollins = hadron.mFrag - mUnpol->Evaluate(parton, hadron, scale);
         
         return std::vector<Hadron>(1, hadron);
      }

      double PositivityKt::Evaluate(const tppmc::pythpart& parton,
                                    const Hadron& hadron,
                                    double scale) {
         if(not Supports(parton)) {
            return mUnpol->Evaluate(parton, hadron, scale);
         } // if
         // Technically if the hadron probability computations went OK, we
         // shouldn't be able to throw anything other than pions, but in case
         // something weird happened return the unpolarised result.
         switch(hadron.Type()) {
            case  211:
            case -211: // Deliberate fall-through
            case  111: // Deliberate fall-through
               break;
            default:
               return mUnpol->Evaluate(parton, hadron, scale);
               break;
         } // switch(h)
         // Get the function with parameters appropriate for this hadron/parton/
         // scale (remember the functions are in bins in Q, not Q2)
         TF2* function = mFunctions.Find(hadron, sqrt(scale));
         assert(function);
         return function->Eval(hadron.z,
                               hadron.PxPyPz().Phi(),
                               hadron.PxPyPz().Pt());
      }

      bool PositivityKt::Supports(const Parton& parton) const {
         // Up and down quarks only
         return parton.Pdg() == 1 or parton.Pdg() == 2;
      }

      std::string PositivityKt::Name() const {
         std::ostringstream oss;
         oss << "DSS + Collins, kT-dependent positivity bound, <pt2> = " <<
         mMeanPt2;
         return oss.str();
      }

      /**
       Returns true if the fragmentation of the input parton type
       is supported by this function.
       */
      bool PositivityKt::Supports(const Hadron& hadron) const {
         return hadron.type not_eq 111 and abs(hadron.type) not_eq 211;
      }

      void PositivityKt::SetHadron(int i) {
         FragmentationFunction::SetHadron(i);
         mUnpol->SetHadron(i);
      }

      void PositivityKt::SetZRange(double lower, double upper) {
         FragmentationFunction::SetZRange(lower, upper);
         mUnpol->SetZRange(lower, upper);
      }

      double PositivityKt::ZMin() const {
         return mUnpol->ZMin();
      }

      double PositivityKt::ZMax() const {
         return mUnpol->ZMax();
      }
      
      double PositivityKt::MinimumScale() const {
         return mUnpol->MinimumScale();
      }

      double PositivityKt::MaximumScale() const {
         return mUnpol->MaximumScale();
      }


      //
      // class FunctionTable
      //


      FunctionTable::FunctionTable()
      : mQ(NULL) {
      }

      FunctionTable::~FunctionTable() {
         // Delete the functors we created for each ROOT function.
         // Note f is a reference to a pointer so we can modify it to NULL
         // after deletion
         BOOST_FOREACH(FunctorNew*& f, mFunctors) {
            if(f) {
               delete f;
               f = NULL;
            } // if
         } // BOOST_FOREACH
      }

      void FunctionTable::Speak() {
         typedef Table::value_type Entry;
         BOOST_FOREACH(Entry& entry, mFunctions) {
            entry.first.Speak();
            for(int i(1); i <= mQ->GetNbins(); ++i) {
               std::cout << "\t(bin Q = " << mQ->GetBinLowEdge(i) << ") " <<
               entry.second.at(i - 1)->GetName() << " " << 
               entry.second.at(i - 1)->GetTitle() << std::endl;
            } // for
         } // BOOST_FOREACH
      }

      TF2* FunctionTable::Find(const Hadron& hadron, double Q) {
         // Determine which Q bin was requested.
         // If Q is outside the range covered, use the first or last bin
         // as appropriate
         int bin = mQ->FindBin(Q);
         if(bin < 1) { // Too low!
            bin = 1;
         } // if
         else if(bin > mQ->GetNbins()) { // Too high!
            bin = mQ->GetNbins();
         }
         bin -= 1; // Different ranges for TAxis [1,N] and vector [0,N)
         typedef Table::value_type Entry;
         BOOST_FOREACH(Entry& entry, mFunctions) {
            // Entry.first is the FragBin object
            // Entry.second is the vector<TF2*> with Q bins
            if(entry.first.Match(hadron)) { // Found the hadron/parton/depol
               return entry.second.at(bin); // Get the TF2* for this Q bin
            } // if
         } // BOOST_FOREACH
         return NULL;
      }

      TF2* FunctionTable::BuildFunction(FragBin bin, double Q) {
         // I would prefer to pass the functor by value, but I *think* that is
         // the source of some random crashing (ROOT does love it's pointers).
         // At least, I haven't found another source for the crash, so until I
         // do this should stay as a pointer.
         // The FunctorNews are collected in a list and deleted in the
         // destructor.
         mFunctors.push_back(new FunctorNew(mUnpol));
         // Despite what the ROOT documentation says, it seems TF2 needs the
         // class name passed as an argument when using a functor.
         // We also need to make sure the function has a unique "ROOT name"
         // otherwise it seems to crash like a m****r f****r, at least with
         // some ROOT versions.
         std::ostringstream oss;
         oss << "positivity_kt_fcn_" << bin.favoured << "_" << bin.depolMin
         << "_" << bin.depolMax << "_" << mQ->FindBin(Q) - 1;
         TF2* f = new TF2(oss.str().c_str(),
                          mFunctors.back(),
                          0.05, 1., -TMath::Pi(), TMath::Pi(), 3,
                          "tppmc::collins::FunctorNew");
         f->SetParameter(FunctorNew::SCALE, Q);
         f->SetParameter(FunctorNew::FAVOURED, bin.favoured);
         f->SetParameter(FunctorNew::POLARISATION, bin.depolMin);
         // Provide a useful title describing the parameters
         oss.str("");
         oss << mUnpol->Name() << " for Q = " << Q
         << " favoured = " << bin.favoured << " polarisation = "
         << bin.depolMin;
         f->SetTitle(oss.str().c_str());
         return f;
      }

      /** \todo What about multiple initialisations? */
      int FunctionTable::Initialise(const Bins* depolbins, const Bins* Qbins) {
         mUnpol = FragFcnPtr(new tppmc::DssFragmentation);
         typedef std::vector<std::pair<double, double> > MinMaxVec;
         // Get the (min, max) of each bin, in both Q and depolarisation
         MinMaxVec depolBinsMinMax = depolbins->MinMax();
         MinMaxVec QBinMinMax = Qbins->MinMax();
         // Initialise the TAxis with the bin edges
         // We use this to convert Q values into bin number
         mQ = new TAxis(Qbins->Number(), &(Qbins->Edges().at(0)));
         typedef MinMaxVec::value_type MinMax;
         // Vector of NULL TF2 pointers to initialise each entry
         TF2Vec prototype(Qbins->Number(), (TF2*)(NULL));
         BOOST_FOREACH(MinMax d, depolBinsMinMax) {
            BOOST_FOREACH(MinMax Q, QBinMinMax) {
               for(int k(0); k < 3; ++k) { // unfavoured, favoured
                  FragBin bin;
                  bin.depolMin = d.first;
                  bin.depolMax = d.second;
                  bin.favoured = k;
                  if(mFunctions.find(bin) == mFunctions.end()) {
                     mFunctions.insert(std::make_pair(bin, prototype));
                  } // if
                  // Subtract 1 from bin in TAxis due to range [1, N]
                  TF2* f = BuildFunction(bin, Q.first);
                  mFunctions[bin].at(mQ->FindBin(Q.first) - 1) = f;
               } // for
            } // for
         } // for
         return 0;
      }


      //
      // class KtDependent
      //


      std::string KtDependent::Find(const Hadron& hadron) const {
         typedef Table::value_type Pair;
         BOOST_FOREACH(Pair p, mTable) {
            if(p.first.Match(hadron)) return p.second;
         }
         return "";
      }
   } // namespace collins
} // namespace tppmc
